/**
* \file DGVFourierTransform.h
*
* \brief This file defines the class for FFTs using only the FFTW library.
* \author Shekhar S. Chandra, 2008
*
* This file is part of DGV Library.
*
* DGV is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* DGV is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with DGV.  If not, see <http://www.gnu.org/licenses/>.
*/
#ifndef DGVFOURIERTRANSFORM_H
#define DGVFOURIERTRANSFORM_H

#include <complex>
#include <fftw3.h>

#include "DGVAliases.h"

/*! \class DGVFourierTransform
    \brief Applies the various different DFT algorithms to Blitz++ constructs.

    The FFTW library is used to implement the DFT algorithms on Blitz++ Arrays.
    The class can only handle upto 1D/2D/3D transforms and not higher.
    All results are un-normalised. Normalisations must be called using the normalise() function.

    Usage Example:
    \code
    DGVFourierTransform<double,2> Fourier; //!< Fourier Transforms
    Array<complex<double>,2> image, kSpace;

    //Load or create image here
    ...

    Fourier.FFT(image,kSpace);
    \endcode
*/
template<typename type, int rank>
class DGV_TEMPLATEDLL DGVFourierTransform
{
public:
    // -------------------------------------------------------
	/*! \name Constructors */
	//@{
    /*! \fn DGVFourierTransform::DGVFourierTransform(void)
	   	\brief Standard Constructor
	   */
    DGVFourierTransform();
    /*! \fn DGFourierTransform::~DGFourierTransform(void)
	   	\brief Standard Destructor
	   */
    virtual ~DGVFourierTransform();
    //@}

    /*! \fn DGVFourierTransform::FFT(Array<complex<type>,rank> &field, Array<complex<type>,rank> &fftOfField)
	   	\brief The Complex valued Fast Fourier Transform (FFT). Uses the FFTW C library.
    */
	void FFT(Array<complex<type>,rank> &field, Array<complex<type>,rank> &fftOfField);
	/*! \fn DGVFourierTransform::iFFT(Array<complex<type>,rank> &fftOfField, Array<complex<type>,rank> &field)
	   	\brief The Complex valued Inverse Fast Fourier Transform (iFFT). Uses the FFTW C library.
    */
	void iFFT(Array<complex<type>,rank> &fftOfField, Array<complex<type>,rank> &field);

	/*! \fn DGVFourierTransform::FFT(Array<type,rank> &field, Array<complex<type>,rank> &fftOfField)
	   	\brief The Real valued Fast Fourier Transform (Real FFT). Uses the FFTW C library.
    */
	void FFT_Real2Complex(Array<type,rank> &field, Array<complex<type>,rank> &fftOfField);
	/*! \fn DGVFourierTransform::FFT(Array<type,rank> &field, Array<complex<type>,rank> &fftOfField)
	   	\brief The Real valued Inverse Fast Fourier Transform in the West (Real iFFT). Uses the FFTW C library.
    */
	void FFT_Complex2Real(Array<complex<type>,rank> &fftOfField, Array<type,rank> &field);

    /*! \fn DGVFourierTransform::normalise(Array<complex<type>,rank> &field)
	   	\brief Normalise the result of the Transform
    */
    void normalise(Array<complex<type>,rank> &field);
    /*! \fn DGVFourierTransform::normalise(Array<type,rank> &field)
	   	\brief Normalise the result of the Transform
    */
    void normalise(Array<type,rank> &field);
    /*! \fn DGVFourierTransform::shiftDC(Array<type,rank> &field)
	   	\brief Shift the DC by half the array size periodically.
    */
    Array<type,rank> shiftDC(Array<type,rank> &field);
    /*! \fn DGVFourierTransform::shiftDC(Array<complex<type>,rank> &field)
	   	\brief Shift the DC by half the array size periodically.
    */
    Array<complex<type>,rank> shiftDC(Array<complex<type>,rank> &field);
	/*! \fn DGVFourierTransform::reset()
	   	\brief Reset the Transform
    */
    inline void reset()
    {   mode = FFTW_ESTIMATE;   fftw_cleanup(); }

protected:
    unsigned mode; //!< FFTW Plan Flag
    fftw_plan plan; //!< FFTW Paln for DFT's

    /*! \fn DGVFourierTransform::execute()
	   	\brief Execute the plan constructed
    */
    void execute();
};

template<typename type, int rank>
DGVFourierTransform<type,rank>::DGVFourierTransform()
{
    mode = FFTW_ESTIMATE;
}

template<typename type, int rank>
DGVFourierTransform<type,rank>::~DGVFourierTransform()
{
    fftw_cleanup();
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::FFT(Array<complex<type>,rank> &field, Array<complex<type>,rank> &fftOfField)
{
    if(rank < 4 && rank > 0) //Protect FFTW
    {
        const int extent[3] = {field.rows(),field.cols(),field.depth()};
        ///\todo Careful: Extent in 3rd Dimension is non zero for 2D arrays
        //cerr << "FFT Extent: " << extent[0] << ", " << extent[1] << ", " << extent[2] << endl;

        plan = fftw_plan_dft(rank,extent,reinterpret_cast<fftw_complex*>(field.data()),reinterpret_cast<fftw_complex*>(fftOfField.data()),FFTW_FORWARD,mode);

        execute();
    }
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::iFFT(Array<complex<type>,rank> &fftOfField, Array<complex<type>,rank> &field)
{
    if(rank < 4 && rank > 0) //Protect FFTW
    {
        const int extent[3] = {fftOfField.rows(),fftOfField.cols(),fftOfField.depth()};
        ///\todo Careful: Extent in 3rd Dimension is non zero for 2D arrays
        //cerr << "iFFT Extent: " << extent[0] << ", " << extent[1] << ", " << extent[2] << endl;

        plan = fftw_plan_dft(rank,extent,reinterpret_cast<fftw_complex*>(fftOfField.data()),reinterpret_cast<fftw_complex*>(field.data()),FFTW_BACKWARD,mode);

        execute();
    }
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::FFT_Real2Complex(Array<type,rank> &field, Array<complex<type>,rank> &fftOfField)
{
    if(rank < 4 && rank > 0) //Protect FFTW
    {
        const int extent[3] = {field.rows(),field.cols(),field.depth()};
        ///\todo Careful: Extent in 3rd Dimension is non zero for 2D arrays
        //cerr << "FFT Extent: " << extent[0] << ", " << extent[1] << ", " << extent[2] << endl;

        plan = fftw_plan_dft_r2c(rank,extent,field.data(),reinterpret_cast<fftw_complex*>(fftOfField.data()),mode);

        execute();
    }
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::FFT_Complex2Real(Array<complex<type>,rank> &fftOfField, Array<type,rank> &field)
{
    if(rank < 4 && rank > 0) //Protect FFTW
    {
        const int extent[3] = {fftOfField.rows(),fftOfField.cols(),fftOfField.depth()};
        ///\todo Careful: Extent in 3rd Dimension is non zero for 2D arrays
        //cerr << "FFT Extent: " << extent[0] << ", " << extent[1] << ", " << extent[2] << endl;

        plan = fftw_plan_dft_c2r(rank,extent,reinterpret_cast<fftw_complex*>(fftOfField.data()),field.data(),mode);

        execute();
    }
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::normalise(Array<complex<type>,rank> &field)
{
    type factor;

    if(rank == 1)
        factor = static_cast<type>(field.size());
    else if(rank == 2)
        factor = static_cast<type>(field.rows()*field.cols());
    else
        factor = static_cast<type>(field.rows()*field.cols()*field.depth());

    //cerr << "Normalising by " << factor << endl;
    field /= factor;
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::normalise(Array<type,rank> &field)
{
    type factor;

    if(rank == 1)
        factor = static_cast<type>(field.size());
    else if(rank == 2)
        factor = static_cast<type>(field.rows()*field.cols());
    else
        factor = static_cast<type>(field.rows()*field.cols()*field.depth());

    //cerr << "Normalising by " << factor << endl;
    field /= factor;
}

template<typename type, int rank>
Array<type,rank> DGVFourierTransform<type,rank>::shiftDC(Array<type,rank> &field)
{
    int rowsHalf = static_cast<int>(field.rows()/2.0), colsHalf = static_cast<int>(field.cols()/2.0);
    Array<type,rank> tmpField(field.shape());

    for(int j = 0; j < field.rows(); j ++)
        for(int k = 0; k < field.cols(); k ++)
            tmpField( (j+rowsHalf)%field.rows() , (k+colsHalf)%field.cols() ) = field(j,k);

    return tmpField;
}

template<typename type, int rank>
Array<complex<type>,rank> DGVFourierTransform<type,rank>::shiftDC(Array<complex<type>,rank> &field)
{
    int rowsHalf = static_cast<int>(field.rows()/2.0), colsHalf = static_cast<int>(field.cols()/2.0);
    Array<complex<type>,rank> tmpField(field.shape());

    for(int j = 0; j < field.rows(); j ++)
        for(int k = 0; k < field.cols(); k ++)
            tmpField( (j+rowsHalf)%field.rows() , (k+colsHalf)%field.cols() ) = field(j,k);

    return tmpField;
}

//Private
template<typename type, int rank>
void DGVFourierTransform<type,rank>::execute()
{
    if(plan != NULL)
    {
        fftw_execute(plan);
        fftw_destroy_plan(plan);
    }
}

#endif // DGVFOURIERTRANSFORM_H
